#include "Tree.hpp"
#include <stack>
namespace broadphase{

void Tree::computeDivisions(std::vector<unsigned int>& axisStack)
{
    vec3 dimensions(_bb.xMax() - _bb.xMin(), _bb.yMax() - _bb.yMin(), _bb.zMax() - _bb.zMin());

    axisStack.reserve(32);

    for (unsigned int level = 0; level < 32; ++level)
    {
        int axis = 0;
        if (dimensions[0] >= dimensions[1])
        {
            if (dimensions[0] >= dimensions[2]) axis = 0;
            else axis = 2;
        }
        else if (dimensions[1] >= dimensions[2]) axis = 1;
        else axis = 2;

        axisStack.push_back(axis);
        dimensions[axis] /= 2.0f;
    }
}

int Tree::divide(AABB& bb, int nodeIndex, unsigned int level,
    const std::vector<vec3>& centers,
    const std::vector<uint>& axisStack)
{
    const static scalar epsilon = 1e-6;

    Node& node = _nodes[nodeIndex];

    bool needToDivide = level < axisStack.size() && (node.first<0 && static_cast<uint>(node.second)>_numPerLeaf);

    if (!needToDivide)
    {
        if (node.first<0)
        {
            int istart = -node.first - 1;
            int iend = istart + node.second;

            // leaf is done, now compute bound on it.
            node.bb.init();
            for (int i = istart; i < iend; ++i)
            {
                node.bb.expandBy(_bbs[_primitives[i]]);
            }

            if (node.bb.valid())
            {
                node.bb._min[0] -= epsilon;
                node.bb._min[1] -= epsilon;
                node.bb._min[2] -= epsilon;
                node.bb._max[0] += epsilon;
                node.bb._max[1] += epsilon;
                node.bb._max[2] += epsilon;
            }
        }

        return nodeIndex;
    }

    int axis = axisStack[level];
    if (node.first<0)
    {
        // leaf node as first <= 0, so look at dividing it.
        // element = -index - 1  ==> index = -element - 1

        int istart = -node.first - 1;
        int iend = istart + node.second;

        scalar original_min = bb._min[axis];
        scalar original_max = bb._max[axis];

        scalar mid = (original_min + original_max)*0.5f;

        int originalLeftChildIndex = 0;
        int originalRightChildIndex = 0;
        bool insitueDivision = false;

        {
            int left = istart;
            int right = iend - 1;

            while (left<right)
            {
                while (left<right && (centers[_primitives[left]][axis] <= mid)) { ++left; }

                while (left<right && (centers[_primitives[right]][axis]>mid)) { --right; }

                if (left<right)
                {
                    std::swap(_primitives[left], _primitives[right]);
                    ++left;
                    --right;
                }
            }

            if (left == right)
            {
                if (centers[_primitives[left]][axis] <= mid) ++left;
                else --right;
            }

            Node leftLeaf(-istart - 1, (right - istart) + 1);
            Node rightLeaf(-left - 1, iend - left);

            if (leftLeaf.second <= 0)
            {
                originalLeftChildIndex = 0;
                originalRightChildIndex = nodeIndex;
                insitueDivision = true;
            }
            else if (rightLeaf.second <= 0)
            {
                originalLeftChildIndex = nodeIndex;
                originalRightChildIndex = 0;
                insitueDivision = true;
            }
            else
            {
                originalLeftChildIndex = addNode(leftLeaf);
                originalRightChildIndex = addNode(rightLeaf);
            }
        }

        scalar restore = bb._max[axis];
        bb._max[axis] = mid;
        int leftChildIndex = originalLeftChildIndex != 0 ? divide(bb, originalLeftChildIndex, level + 1, centers, axisStack) : 0;
        bb._max[axis] = restore;

        restore = bb._min[axis];
        bb._min[axis] = mid;
        int rightChildIndex = originalRightChildIndex != 0 ? divide(bb, originalRightChildIndex, level + 1, centers, axisStack) : 0;
        bb._min[axis] = restore;

        if (!insitueDivision)
        {
            // take a second reference to node we are working on as the std::vector<> resize could
            // have invalidate the previous node ref.
            Node& newNodeRef = _nodes[nodeIndex];

            newNodeRef.first  = leftChildIndex;
            newNodeRef.second = rightChildIndex;

            newNodeRef.bb.init();
            if (leftChildIndex != 0)  newNodeRef.bb.expandBy(_nodes[leftChildIndex].bb);
            if (rightChildIndex != 0) newNodeRef.bb.expandBy(_nodes[rightChildIndex].bb);

            if (!newNodeRef.bb.valid())
            {
                assert(false);
            }
        }
    }
    else
    {
        assert(false);
    }

    return nodeIndex;
}

bool Tree::build(const std::vector<AABB>& bbs)
{
    if (bbs.size() < _numPerLeaf) return false;
    
    size_t count = bbs.size();
    _bbs = bbs;

    _bb.init();
    for (size_t i = 0; i < count; i++)
        _bb.expandBy(_bbs[i]);

    _primitives.resize(count);
    for (uint i = 0; i < count; ++i)
        _primitives[i] = i;

    std::vector<uint> axisStack;
    computeDivisions(axisStack);

    std::vector<vec3> centers;
    centers.reserve(count);
    for (const auto& bb : _bbs)
        centers.push_back(bb.center());
    
    _nodes.reserve(2 * count);

    // root
    Node node(-1, count);
    node.bb = _bb;

    int nodeNum = addNode(node);
    AABB bb = _bb;
    divide(bb, nodeNum, 0, centers, axisStack);

    _nodes.shrink_to_fit();

    // ordered

    return !_nodes.empty();
}

void Tree::queryLeaf(uint leaf, CollideResults& results)
{
    const Node& node = _nodes[leaf];
    int istart = -node.first - 1;
    int iend = istart + node.second;
    
    for (int i = istart; i < iend; ++i)
    {
        for (int j = i + 1; j < iend; ++j)
        {
            if (_bbs[_primitives[i]].overlap(_bbs[_primitives[j]]))
            {
                results.push_back(std::make_pair(_primitives[i], _primitives[j]));
            }
        }
    }
}

void Tree::crossChildren(uint n, CollideResults& results)
{
    if (!_nodeCrossed[n])
    {
        if (_nodes[n].isLeaf())
            queryLeaf(n, results);
        else 
            queryNode(_nodes[n].first, _nodes[n].second, results);
        _nodeCrossed[n] = true;
    }
}

void Tree::queryNode(uint l, uint r, CollideResults& results)
{
    crossChildren(l, results);
    crossChildren(r, results);
    
    const Node& n0 = _nodes[l];
    const Node& n1 = _nodes[r];

    if (!n1.bb.overlap(n0.bb))
        return;
    
    if (n0.isLeaf())
    {
        if (n1.isLeaf())
        {
            // query leaves
            int ls = -n0.first - 1;
            int le = ls + n0.second;
            int rs = -n1.first - 1;
            int re = rs + n1.second;

            for (int i = ls; i < le; ++i)
                for (int j = rs; j < re; ++j)
                    if (_bbs[_primitives[i]].overlap(_bbs[_primitives[j]]))
                        results.push_back(std::make_pair(_primitives[i], _primitives[j]));
        }
        else
        {
            // n0 is leaf, n1 is not
            queryNode(l, n1.first,  results);
            queryNode(l, n1.second, results);
        }
    }
    else
    {
        if (n1.isLeaf())
        {
            // n1 is leaf, n0 is not
            queryNode(n0.first,  r, results);
            queryNode(n0.second, r, results);
        }
        else
        {
            queryNode(n0.first,  n1.first,  results);
            queryNode(n0.first,  n1.second, results);
            queryNode(n0.second, n1.first,  results);
            queryNode(n0.second, n1.second, results);
        }
    }
}

CollideResults Tree::query()
{
    CollideResults results;
    if (_bbs.size() > 1) 
    {
        if (_numPerLeaf == 1) 
        {
            if (!_nodes[0].isLeaf()) 
            {
                _nodeCrossed.resize(_nodes.size(), false);
                queryNodeSingle(_nodes[0].first, _nodes[0].second, results);
            }
        }
        else 
        {
            if (!_nodes[0].isLeaf())
            {
                _nodeCrossed.resize(_nodes.size(), false);
                queryNode(_nodes[0].first, _nodes[0].second, results);
            }
            else
            {
                queryLeaf(0, results);
            }
        }
    }

    return results;
}

std::vector<uint> Tree::query(const AABB& aabb)
{
    std::vector<uint> results;
    
    std::stack<uint> stack;
    stack.push(0);
    while (stack.size() > 0) 
    {
        uint index = stack.top();
        stack.pop();
        const Node& node = _nodes[index];
        if (node.bb.overlap(aabb)) 
        {
            if (node.isLeaf()) 
            {
                int istart = -node.first - 1;
                int iend = istart + node.second;
                for (int i = istart; i < iend; ++i) 
                {
                    if (_bbs[_primitives[i]].overlap(aabb))
                        results.push_back(_primitives[i]);
                }
            }
            else 
            {
                stack.push(node.first);
                stack.push(node.second);
            }
        }
    }

    return results;
}

void Tree::crossChildrenSingle(uint n, CollideResults& results)
{
    if (!_nodeCrossed[n])
    {
        queryNodeSingle(_nodes[n].first, _nodes[n].second, results);
        _nodeCrossed[n] = true;
    }
}

void Tree::queryNodeSingle(uint l, uint r, CollideResults& results)
{
    const Node& n0 = _nodes[l];
    const Node& n1 = _nodes[r];

    if (n0.isLeaf())
    {
        if (n1.isLeaf())
        {
            // query leaves
            int i = -n0.first - 1;
            int j = -n1.first - 1;
            if (_bbs[_primitives[i]].overlap(_bbs[_primitives[j]]))
                results.push_back(std::make_pair(_primitives[i], _primitives[j]));
        }
        else
        {
            crossChildrenSingle(r, results);
            if (!n1.bb.overlap(n0.bb)) return;
            queryNodeSingle(l, n1.first, results);
            queryNodeSingle(l, n1.second, results);
        }
    }
    else
    {
        if (n1.isLeaf())
        {
            crossChildrenSingle(l, results);
            if (!n1.bb.overlap(n0.bb)) return;
            queryNodeSingle(n0.first, r, results);
            queryNodeSingle(n0.second, r, results);
        }
        else
        {
            crossChildrenSingle(l, results);
            crossChildrenSingle(r, results);
            if (!n1.bb.overlap(n0.bb)) return;
            queryNodeSingle(n0.first, n1.first, results);
            queryNodeSingle(n0.first, n1.second, results);
            queryNodeSingle(n0.second, n1.first, results);
            queryNodeSingle(n0.second, n1.second, results);
        }
    }
}

}